from pydantic import BaseModel
from typing import List, Optional

class QueryRequest(BaseModel):
    question: str
    top_k: Optional[int] = 5
    max_tokens: Optional[int] = 1000
    temperature: Optional[float] = 0.7

class QueryResponse(BaseModel):
    question: str
    answer: str
    retrieved_context: List[dict]
    context_count: int

class DocumentUploadResponse(BaseModel):
    document_external_id: str 
    document_title: Optional[str] = None
    total_chunks: Optional[int] = None
    file_size: Optional[int] = None
    status: str
